import json

import requests

from config import API_KEY, SECRET_KEY


# 获取 access_token 的函数
def get_access_token():
    url = "https://aip.baidubce.com/oauth/2.0/token"
    params = {"grant_type": "client_credentials", "client_id": API_KEY, "client_secret": SECRET_KEY}
    return str(requests.post(url, params=params).json().get("access_token"))


# 大模型选择Agent的函数
def determine_agent(user_input):
    # 优化后的 system_text，支持中英文
    system_text = (
        "You are a role selector capable of choosing the appropriate role based on the user's input. "
        "Carefully analyze the user input to determine whether it relates to travel, such as 'travel plan', 'itinerary', or 'destination recommendation', "
        "or to content creation, such as 'blog', 'social media content', or 'marketing copy'. "
        "Then, reply only with the current role as 'Travel Assistant' or 'Content Creator' without any additional content or explanations. "
        "Ensure your response includes only the role selection and strictly adheres to the user's intent. Absolutely no content outside the predefined responses."
        "Please respond only in English."
    )

    # 在用户输入前加上 system_text
    full_input = system_text + user_input

    url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token={get_access_token()}"

    payload = json.dumps({
        "messages": [{"role": "user", "content": full_input}],
        "temperature": 0.8,
        "top_p": 0.8,
        "penalty_score": 1,
        "disable_search": False,
        "enable_citation": False,
        "collapsed": True,
        "system": system_text
    })

    headers = {'Content-Type': 'application/json'}
    response = requests.post(url, headers=headers, data=payload)
    if response.status_code == 200:
        result = response.json().get("result", "无返回结果")
        print(f"获取的结果: {result}")  # 打印当前获取的结果
        # 提取角色
        if "Travel Assistant" in result:
            return "travel_assistant", result
        elif "Content Creator" in result:
            return "content_creator", result
    return None, "API请求失败，请检查输入或网络连接。"
